# SimpleOutcar.py

import re
from collections import Counter

import numpy as np

from SimplePoscar import Atom, Atoms


class SimpleOutcar:
    nn_dist: float
    e0: float
    atoms: Atoms
    positions: np.ndarray
    forces: np.ndarray

    def __init__(self, filepath: str):
        self.read_outcar(filepath)

    def _get_nn_dist(self, text: str):
        dists = re.findall(r"\d+\s(\d+\.\d+)", text)
        c = Counter(dists)
        dist = c.most_common()[0][0]

        return float(dist)

    def _get_cell(self, text: str) -> np.ndarray:
        lines = text.split("\n")
        idx = 0
        for i, line in enumerate(lines):
            if "direct lattice vectors" in line:
                idx = i + 1
                break

        cell = []
        for i in range(idx, idx + 3):
            line = lines[i]
            a = line.split()[:3]
            cell.append(a)

        return np.array(cell, dtype=float)

    def _get_positions(self, text: str, atom_num: int):
        lines = text.split("\n")
        idx = 0
        for i, line in enumerate(lines):
            if "TOTAL-FORCE" in line:
                idx = i + 2
                break

        positions = []
        forces = []
        for i in range(idx, idx + atom_num):
            x, y, z, fx, fy, fz = [float(p) for p in lines[i].split()]
            positions.append((x, y, z))
            forces.append((fx, fy, fz))

        return np.array(positions), np.array(forces)

    def _get_atoms(self, cell: np.ndarray, counts: list[int], symbols: list[str], positions: np.ndarray):
        atoms = Atoms(cell=cell, is_direct=False)

        pos_idx = 0
        for symbol, count in zip(symbols, counts):
            for i in range(pos_idx, pos_idx + count):
                atoms.append(Atom(index=pos_idx,
                                  symbol=symbol,
                                  coord=positions[i]))
            pos_idx += count

        return atoms

    def read_outcar(self, filepath: str):
        nn_dist = 0.0
        e0 = 0.0
        # atom_num = 0
        counts = []
        symbols = []
        cell = np.zeros((3, 3))
        positions = np.zeros((0, 3))
        forces = np.zeros((0, 3))

        with open(filepath, "r") as f:
            text = f.read()

        # Split text into blocks using the separator
        blocks = re.split(
            r"------------------------------------------------------------------------------------------------",
            text)

        for block in blocks:
            if match := re.search(r"nearest neighbor table", block):
                # nearest neighbor table
                nn_dist = self._get_nn_dist(block)
                break
        # for block in blocks:
        #     if match := re.search(r"NIONS\s*=\s*(\d+)", block):
        #         # 'NIONS =     54'
        #         atom_num = int(match.group(1))
        for block in blocks:
            if match := re.search(r"ions per type =\s*(\d+(?:\s+\d+)*)", block):
                # ions per type =    1  2  3
                counts = list(map(int, re.findall(r"\d+", match.group(1))))
                break
        for block in blocks:
            if matches := re.findall(r"\s*VRHFIN\s*=\s*([A-Z][a-z]?)\s*\:", block):
                symbols = matches
                break
        for block in reversed(blocks):
            if match := re.search(r"energy\(sigma->0\)\s*=\s*(-?\d+\.?\d*)", block):
                # energy(sigma->0) = -1.23
                e0 = float(match.group(1))
                break
        for block in reversed(blocks):
            if match := re.search(r"direct lattice vectors", block):
                # direct lattice vectors
                cell = self._get_cell(block)
                break
        for block in reversed(blocks):
            if match := re.search(r"POSITION\s*TOTAL-FORCE", block):
                # POSITION TOTAL-FORCE (eV/Angst)
                positions, forces = self._get_positions(block, sum(counts))
                break

        # print(nn_dist, e0, sum(counts), counts, symbols, cell, len(positions), positions)

        self.nn_dist = nn_dist
        self.e0 = e0
        self.atoms = self._get_atoms(cell, counts, symbols, positions)
        self.positions = positions
        self.forces = forces
